import re
from pathlib import Path

from pydantic_settings import BaseSettings, SettingsConfigDict

import pymysql

class _Settings(BaseSettings):
    model_config = SettingsConfigDict(
        env_file="../../.env",
        env_file_encoding="utf-8",
        env_ignore_empty=True,
        extra="ignore",
    )

    DB_SERVER: str
    DB_PORT: int
    DB_DB: str
    DB_USER: str
    DB_PASSWORD: str

settings = _Settings() # type: ignore

# Connect to the database
db = pymysql.connect(
    host=settings.DB_SERVER,
    port=settings.DB_PORT,
    user=settings.DB_USER,
    password=settings.DB_PASSWORD,
    db=settings.DB_DB,
)

def _query_db(query):
    cursor = db.cursor()
    cursor.execute(query)
    result = cursor.fetchall()
    return result

def _get_tables(tables: list[str] = None):
    query = f"""
        SELECT
          table_name,
          table_comment
        FROM
          information_schema.TABLES
        WHERE
          table_schema = '{settings.DB_DB}'
          {f"and table_name in ({",".join([f"'{table}'" for table in tables])})" if tables else ""}
    """
    tables = _query_db(query)
    return tables

def _get_colums(table_name: str):
    query = f"""
        SELECT
          column_name,
          data_type,
          column_type,
          column_key,
          column_comment,
          is_nullable
        FROM
          information_schema.COLUMNS
        WHERE
          table_schema = '{settings.DB_DB}'
          AND table_name = '{table_name}'
        ORDER BY
          ordinal_position ASC;
    """
    colums = _query_db(query)
    return colums

def _table2sqlmodel(generate_dir: Path, tables: list[str] = None):
    tables = _get_tables(tables)

    # Print the rows
    for table_row in tables:
        table_name, table_comment = table_row
        colums = _get_colums(table_name)
        model_str = "from sqlmodel import SQLModel, Field\n\n\n"
        model_str += f"class {_snake2camel(table_name)}(SQLModel, table=True):{f"  # {table_comment}" if table_comment else ""}\n"
        model_str += f"    __tablename__ = \"{table_name}\"\n\n"
        built_in_import_str = set()
        for column_row in colums:
            column_name, data_type, column_type, column_key, column_comment, is_nullable = column_row
            python_data_type = "str"
            data_length = None
            if data_type in ["bigint", "int"]:
                python_data_type = "int"
                if re.match(r"^(big)?int\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["varchar", "text", "longtext"]:
                python_data_type = "str"
                if re.match(r"^varchar\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["datetime", "timestamp"]:
                python_data_type = "datetime"
                built_in_import_str.add("from datetime import datetime")
            elif data_type in ["decimal", "float"]:
                python_data_type = "float"
            elif data_type in ["date"]:
                python_data_type = "date"
                built_in_import_str.add("from datetime import date")
            elif data_type in ["time"]:
                python_data_type = "time"
                built_in_import_str.add("from datetime import time")
            model_str += f"    {column_name}: {python_data_type} | None = Field({"primary_key=True, " if column_key == "PRI" else "index=True, " if column_key else ""}{"default=None, " if is_nullable == "YES" else ""}{f"max_length={data_length}, " if data_length else ""}){f"  # {column_comment}" if column_comment else ""}\n"

        if built_in_import_str:
            model_str =  "\n".join(built_in_import_str) + "\n\n" + model_str
        print(model_str + "\n")
        with open(generate_dir.joinpath(f"{table_name}.py"), "w", encoding="utf-8") as f:
            f.write(model_str)

def _table2sqlalchemy(generate_dir: Path, tables: list[str] = None):
    tables = _get_tables(tables)

    # Print the rows
    for table_row in tables:
        table_name, table_comment = table_row
        colums = _get_colums(table_name)
        model_str = "from sqlalchemy.orm import mapped_column, Mapped\n\nfrom app.models.base import Base\n\n\n"
        model_str += f"class {_snake2camel(table_name)}(Base):{f"  # {table_comment}" if table_comment else ""}\n"
        model_str += f"    __tablename__ = \"{table_name}\"\n\n"
        built_in_import_str = set()
        sqlalchemy_built_in_import_str = set()
        for column_row in colums:
            column_name, data_type, column_type, column_key, column_comment, is_nullable = column_row
            python_data_type = "str"
            sqlalchemy_data_type = "str"
            data_length = None
            if data_type in ["bigint", "int"]:
                python_data_type = "int"
                sqlalchemy_data_type = "Integer"
                if re.match(r"^(big)?int\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["varchar", "text", "longtext"]:
                python_data_type = "str"
                sqlalchemy_data_type = "String"
                if re.match(r"^varchar\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["datetime", "timestamp"]:
                python_data_type = "datetime"
                sqlalchemy_data_type = "DateTime"
                built_in_import_str.add("from datetime import datetime")
            elif data_type in ["decimal", "float"]:
                python_data_type = "float"
                sqlalchemy_data_type = "Float"
            elif data_type in ["date"]:
                python_data_type = "date"
                sqlalchemy_data_type = "Date"
                built_in_import_str.add("from datetime import date")
            elif data_type in ["time"]:
                python_data_type = "time"
                sqlalchemy_data_type = "Time"
                built_in_import_str.add("from datetime import time")
            sqlalchemy_built_in_import_str.add(sqlalchemy_data_type)
            model_str += f"    {column_name}: Mapped[{python_data_type}] = mapped_column({sqlalchemy_data_type}{f"({data_length})" if data_length and python_data_type != 'int' else ""}"
            model_str += f"{', primary_key=True' if column_key == 'PRI' else ''}"
            model_str += f"{', unique=True' if column_key == 'UNI' else ''}"
            model_str += f"{', index=True' if column_key == 'MUL' else ''}"
            model_str += f"{f', comment=\"{column_comment}\"' if column_comment else ''})\n"

        if sqlalchemy_built_in_import_str:
            model_str = "from sqlalchemy import " + ", ".join(sqlalchemy_built_in_import_str) + "\n" + model_str
        if built_in_import_str:
            model_str =  "\n".join(built_in_import_str) + "\n\n" + model_str
        print(model_str + "\n")
        with open(generate_dir.joinpath(f"{table_name}.py"), "w", encoding="utf-8") as f:
            f.write(model_str)

def _table2basemodel(generate_dir: Path, tables: list[str] = None):
    tables = _get_tables(tables)

    # Print the rows
    for table_row in tables:
        table_name, table_comment = table_row
        colums = _get_colums(table_name)
        model_str = f"class {_snake2camel(table_name)}DTO(BaseModel):{f"  # {table_comment}" if table_comment else ""}\n"
        built_in_import_str = set()
        for column_row in colums:
            column_name, data_type, column_type, column_key, column_comment, is_nullable = column_row
            python_data_type = "str"
            data_length = None
            if data_type in ["bigint", "int"]:
                python_data_type = "int"
                if re.match(r"^(big)?int\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["varchar", "text", "longtext"]:
                python_data_type = "str"
                if re.match(r"^varchar\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["datetime", "timestamp"]:
                python_data_type = "datetime"
                built_in_import_str.add("from datetime import datetime")
            elif data_type in ["decimal", "float"]:
                python_data_type = "float"
            elif data_type in ["date"]:
                python_data_type = "date"
                built_in_import_str.add("from datetime import date")
            elif data_type in ["time"]:
                python_data_type = "time"
                built_in_import_str.add("from datetime import time")
            model_str += f"    {column_name}: {python_data_type} | None = Field({"default=None, " if is_nullable == "YES" else ""}{f"max_length={data_length}, " if data_length and python_data_type != 'int' else ""}{f"description=\"{column_comment}\", " if column_comment else ""})\n"

        model_str = "from pydantic import Field\n\nfrom app.schemas.base import BaseModel\n\n\n" + model_str
        if built_in_import_str:
            model_str = "\n".join(built_in_import_str) + "\n\n" + model_str
        print(model_str + "\n")
        with open(generate_dir.joinpath(f"{table_name}.py"), "w", encoding="utf-8") as f:
            f.write(model_str)

def _table2basemodel_allinone(generate_dir: Path, tables: list[str] = None):
    tables = _get_tables(tables)

    whole_model_str = ""
    whole_built_in_import_str = set()
    # Print the rows
    for table_row in tables:
        table_name, table_comment = table_row
        colums = _get_colums(table_name)
        model_str = f"class {_snake2camel(table_name)}DTO(BaseModel):{f"  # {table_comment}" if table_comment else ""}\n"
        built_in_import_str = set()
        for column_row in colums:
            column_name, data_type, column_type, column_key, column_comment, is_nullable = column_row
            python_data_type = "str"
            data_length = None
            if data_type in ["bigint", "int"]:
                python_data_type = "int"
                if re.match(r"^(big)?int\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["varchar", "text", "longtext"]:
                python_data_type = "str"
                if re.match(r"^varchar\(\d+\)", column_type):
                    data_length = int(re.findall(r"\d+", column_type)[0])
            elif data_type in ["datetime", "timestamp"]:
                python_data_type = "datetime"
                built_in_import_str.add("from datetime import datetime")
            elif data_type in ["decimal", "float"]:
                python_data_type = "float"
            elif data_type in ["date"]:
                python_data_type = "date"
                built_in_import_str.add("from datetime import date")
            elif data_type in ["time"]:
                python_data_type = "time"
                built_in_import_str.add("from datetime import time")
            model_str += f"    {column_name}: {python_data_type} | None = Field({"default=None, " if is_nullable == "YES" else ""}{f"max_length={data_length}, " if data_length and python_data_type != 'int' else ""}{f"description=\"{column_comment}\", " if column_comment else ""})\n"
            whole_built_in_import_str |= built_in_import_str

        whole_model_str += model_str + "\n\n"
        model_str = "from pydantic import Field\n\nfrom app.schemas.base import BaseModel\n\n\n" + model_str
        if built_in_import_str:
            model_str = "\n".join(built_in_import_str) + "\n\n" + model_str
        print(model_str + "\n")
    whole_model_str = "from pydantic import Field\n\nfrom app.schemas.base import BaseModel\n\n\n" + whole_model_str
    if whole_built_in_import_str:
        whole_model_str = "\n".join(whole_built_in_import_str) + "\n\n" + whole_model_str
    # Write the whole model to a file
    with open(generate_dir.joinpath("database_schemas.py"), "w", encoding="utf-8") as f:
        f.write(whole_model_str)

def _snake2camel(string):
  return ''.join([_.capitalize() for _ in string.split('_')])

if __name__ == '__main__':
    tables = []
    # _table2sqlmodel(model_dir, tables=tables)
    _table2sqlalchemy(Path.cwd().parent.joinpath("models"), tables=tables)
    _table2basemodel(Path.cwd().parent.joinpath("schemas"), tables=tables)
